#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# @Author: renjin@bit.edu.cn
# @Date  : 2024-08-06

"""
【节点名称】：
    COCOEvalNode
【依赖项安装】：
    pip install spirems
    pip install pycocotools
【订阅类型】：
    无
【发布类型】：
    spirecv_msgs::EvaluationResult （算法评估结果）
    std_msgs::Boolean （发送任务结束标注位）
【构造参数说明】：
    parameter_file (str): 全局参数文件
    sms_shutdown_emit (bool): 是否发送任务结束标注位，默认False
【节点参数】：
    /work_dir (str): 全局工作空间
    /inputs_fn (str): 全局评估任务输入文件夹
    results_dir (str): 评估结果存储文件夹
    supported_datasets (list): 目前可支持的数据 {数据集名称：数据COCO标注文件路径}
【备注】：
    无
"""

import os
import json
import time
import threading
import shutil
import argparse
from typing import Union
from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval
from collections import defaultdict
from spirems import Subscriber, Publisher, cvimg2sms, def_msg, QoS, BaseNode, get_extra_args


def load_dataset(param_spire_coco_gt):
    annotation_file = param_spire_coco_gt
    tic = time.time()
    with open(annotation_file, 'r') as f:
        dataset = json.load(f)
    assert isinstance(dataset, dict), 'annotation file format {} not supported'.format(type(dataset))
    print('Done (t={:0.2f}s)'.format(time.time() - tic))
    print('creating index...')
    anns, cats, imgs = {}, {}, {}
    imgToAnns, catToImgs = defaultdict(list), defaultdict(list)
    if 'annotations' in dataset:
        for ann in dataset['annotations']:
            imgToAnns[ann['image_id']].append(ann)
            anns[ann['id']] = ann

    if 'images' in dataset:
        for img in dataset['images']:
            imgs[img['id']] = img

    if 'categories' in dataset:
        for cat in dataset['categories']:
            cats[cat['id']] = cat

    if 'annotations' in dataset and 'categories' in dataset:
        for ann in dataset['annotations']:
            catToImgs[ann['category_id']].append(ann['image_id'])

    print('index created!')
    _category_name_to_cat_id = {v['name'].replace(" ", "_"): v['id'] for k, v in cats.items()}
    _file_name_to_img_id = {os.path.basename(v['file_name']): v['id'] for k, v in imgs.items()}
    return _category_name_to_cat_id, _file_name_to_img_id


class COCOEvalNode(threading.Thread, BaseNode):
    def __init__(
        self,
        job_name: str,
        ip: str = '127.0.0.1',
        port: int = 9094,
        param_dict_or_file: Union[dict, str] = None,
        sms_shutdown: bool = True,
        **kwargs
    ):
        threading.Thread.__init__(self)
        sms_shutdown = True if sms_shutdown in ['True', 'true', '1', True] else False
        BaseNode.__init__(
            self,
            self.__class__.__name__,
            job_name,
            ip=ip,
            port=port,
            param_dict_or_file=param_dict_or_file,
            sms_shutdown=sms_shutdown,
            **kwargs
        )
        print("  \033[32m{}\033[0m = {}".format("sms_shutdown", sms_shutdown))
        self.sms_shutdown_emit = self.get_param("sms_shutdown_emit", False)
        self.remote_ip = self.get_param("remote_ip", "127.0.0.1")
        self.remote_port = self.get_param("remote_port", 9094)
        self.g_work_dir = self.get_param("/work_dir", "C:/tmp")
        self.g_inputs_fn = self.get_param("/inputs_fn", 'COCOEvalNodeInputs')
        self.results_fn = self.get_param("results_dir", 'COCOEvalNodeResults')
        self.supported_datasets = self.get_param(
            "supported_datasets",
            {"coco_detection": "D:/dataset/annotations/instances_val2017.json"}
        )
        self.params_help()

        if self.sms_shutdown_emit:
            self._shutdown_writer = Publisher(
                '/' + job_name + '/shutdown', 'std_msgs::Boolean',
                ip=ip, port=port, qos=QoS.Reliability
            )

        self.start()

    def release(self):
        BaseNode.release(self)
        if self.sms_shutdown_emit:
            self._shutdown_writer.kill()
            self._shutdown_writer.join()

    def run(self):
        inputs_dir = os.path.join(self.g_work_dir, self.g_inputs_fn)
        results_dir = os.path.join(self.g_work_dir, self.results_fn)
        if os.path.exists(inputs_dir):
            shutil.rmtree(inputs_dir)
        os.makedirs(inputs_dir)
        if not os.path.exists(results_dir):
            os.makedirs(results_dir)

        while self.is_running():
            time.sleep(1)

            items_in_dir = os.listdir(str(inputs_dir))
            for item in items_in_dir:
                if item.endswith('.json'):
                    time.sleep(1)
                    file_path = os.path.join(str(inputs_dir), item)
                    with open(file_path, 'r') as f:
                        data = json.load(f)
                        spire_coco_gt = self.supported_datasets[data['dataset']]
                        spire_coco_dt_dir = data['coco_dt_dir']
                        fps = data['fps']

                    category_name_to_cat_id, file_name_to_img_id = load_dataset(spire_coco_gt)
                    assembled_results = []
                    for json_file in os.listdir(spire_coco_dt_dir):
                        if os.path.splitext(json_file)[-1] == '.json':
                            with open(os.path.join(spire_coco_dt_dir, json_file), 'r') as f2:
                                json_dict = json.load(f2)
                            boxes, scores, labels = [], [], []
                            for anno in json_dict['targets']:
                                boxes.append(anno['bbox'])
                                scores.append(anno['score'])
                                labels.append(category_name_to_cat_id[anno['category_name']])
                            assembled_results.extend(
                                [
                                    {
                                        "image_id": file_name_to_img_id[json_dict['file_name']],
                                        "category_id": labels[k],
                                        "bbox": box,
                                        "score": scores[k],
                                    }
                                    for k, box in enumerate(boxes)
                                ]
                            )
                    spire_coco_dt = os.path.join(self.g_work_dir, 'detections.json')
                    with open(spire_coco_dt, "w") as f:
                        json.dump(assembled_results, f)

                    coco_gt = COCO(spire_coco_gt)
                    coco_dt = coco_gt.loadRes(str(spire_coco_dt))
                    coco_eval = COCOeval(coco_gt, coco_dt, "bbox")
                    coco_eval.evaluate()
                    coco_eval.accumulate()
                    coco_eval.summarize()

                    msg = def_msg('spirecv_msgs::EvaluationResult')
                    msg['client_id'] = os.path.splitext(item)[0]
                    msg['metrics'] = ['AP', 'AP50', 'AP75', 'APs', 'APm', 'APl', 'FPS']
                    msg['results'] = coco_eval.stats[:6].tolist()
                    msg['results'].append(fps)
                    metric_writer = Publisher(
                        '/SpireView/{}/TaskResults'.format(data['client_id']), 'spirecv_msgs::EvaluationResult',
                        ip=self.remote_ip, port=self.remote_port
                    )
                    metric_writer.publish(msg)
                    metric_writer.kill()
                    metric_writer.join()

                    with open(os.path.join(str(results_dir), item), 'w') as f:
                        json.dump(msg, f, indent=4)

                    os.remove(file_path)

                    if self.sms_shutdown_emit:
                        msg = def_msg('std_msgs::Boolean')
                        msg['data'] = True
                        self._shutdown_writer.publish(msg)

        self.release()
        print('{} quit!'.format(self.__class__.__name__))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--config',
        type=str,
        default='default_params.json',
        help='SpireCV2 Config (.json)')
    parser.add_argument(
        '--job-name',
        type=str,
        default='live',
        help='SpireCV Job Name')
    parser.add_argument(
        '--ip',
        type=str,
        default='127.0.0.1',
        help='SpireMS Core IP')
    parser.add_argument(
        '--port',
        type=int,
        default=9094,
        help='SpireMS Core Port')
    # args = parser.parse_args()
    args, unknown_args = parser.parse_known_args()
    if not os.path.isabs(args.config):
        current_path = os.path.abspath(__file__)
        params_dir = os.path.join(current_path[:current_path.find('spirecv-pro') + 11], 'params', 'spirecv2')
        args.config = os.path.join(params_dir, args.config)
    print("--config:", args.config)
    print("--job-name:", args.job_name)
    extra = get_extra_args(unknown_args)

    node = COCOEvalNode(args.job_name, param_dict_or_file=args.config, ip=args.ip, port=args.port, **extra)
    node.spin()
